/*
 * Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
 * with the License. A copy of the License is located at
 *
 * http://aws.amazon.com/apache2.0/
 *
 * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
 * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
 * and limitations under the License.
 */
package com.yhzdys.ocr.service.support;

import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.modality.cv.output.BoundingBox;
import ai.djl.modality.cv.output.DetectedObjects;
import ai.djl.modality.cv.output.Rectangle;
import ai.djl.modality.cv.util.NDImageUtils;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.NDManager;
import ai.djl.translate.NoBatchifyTranslator;
import ai.djl.translate.Translator;
import ai.djl.translate.TranslatorContext;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * A {@link Translator} that post-process the {@link NDArray} into {@link DetectedObjects} with
 * boundaries.
 */
public class PpWordDetectionTranslator implements NoBatchifyTranslator<Image, DetectedObjects> {

    private static final int maxLength = 1280;
    private final float[] mean = new float[]{0.485f, 0.456f, 0.406f};
    private final float[] std = new float[]{0.229f, 0.224f, 0.225f};

    /**
     * {@inheritDoc}
     */
    @Override
    public NDList processInput(TranslatorContext ctx, Image input) {
        NDManager ctxNDManager = ctx.getNDManager();
        NDArray image = input.toNDArray(ctxNDManager);

        int h = input.getHeight();
        int w = input.getWidth();
        int[] hw = this.scale(h, w, maxLength);

        NDArray resize = NDImageUtils.resize(image, hw[1], hw[0]);
        NDArray tensor = NDImageUtils.toTensor(resize);
        NDArray normalize = NDImageUtils.normalize(tensor, mean, std);
        NDArray expanded = normalize.expandDims(0);
        return new NDList(expanded);
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public DetectedObjects processOutput(TranslatorContext ctx, NDList list) {
        NDArray result = list.singletonOrThrow();
        NDArray squeezed = result.squeeze(0);
        Image image = ImageFactory.getInstance().fromNDArray(squeezed);
        List<BoundingBox> boundingBoxes = image.findBoundingBoxes().stream().filter(box -> {
            Rectangle rect = (Rectangle) box;
            return rect.getWidth() * image.getWidth() > 5 || rect.getHeight() * image.getHeight() > 5;
        }).collect(Collectors.toList());
        OcrUtils.releaseImage(image);

        // 按照从左到右从上到下排序一下
        boundingBoxes.sort((a, b) -> {
            double yDiff = a.getBounds().getY() - b.getBounds().getY();
            // 允许0.002的误差
            if (yDiff <= -0.002D) {
                return -1;
            } else if (yDiff >= 0.002D) {
                return 1;
            } else {
                return Double.compare(a.getBounds().getX(), b.getBounds().getX());
            }
        });

        List<String> names = new ArrayList<>();
        List<Double> probabilities = new ArrayList<>();
        for (int i = 0; i < boundingBoxes.size(); i++) {
            names.add("word");
            probabilities.add(1.0);
        }
        return new DetectedObjects(names, probabilities, boundingBoxes);
    }

    private int[] scale(int h, int w, int max) {
        int localMax = Math.max(h, w);
        float scale = 1.0f;
        if (max < localMax) {
            scale = max * 1.0f / localMax;
        }
        // paddle model only take 32-based size
        return OcrUtils.resize32(h * scale, w * scale);
    }
}
